from utils import *
from model import *
from config import *
from torch.utils import data
import torch

if __name__ =='__main__':
    device = 'cuda' if torch.cuda.is_available() else "cpu"
    dataset = Dataset()
    loader = data.DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        collate_fn=collate_fn
    )

    model = Model()
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)

    for e in range(EPOCH):
        for b, (input, target, mask) in enumerate(loader):
            input, target, mask = input.to(device), target.to(device), mask.to(device)

            loss = model.loss_fn(input, target, mask)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()


        print('>>epoch', e, 'loss', loss.item())
